#ifdef ENABLE_MAPLE_SAN

#include <algorithm>

#include "asan_phases.h"

namespace maple {

// we need to handle SSANode
AddrofNode *MEDoNeighborOpt::getAddrofNode(BaseNode *baseNode) {
  SSANode *ssaCheck = dynamic_cast<SSANode *>(baseNode);
  if (ssaCheck != nullptr)
    return dynamic_cast<AddrofNode *>(dynamic_cast<AddrofSSANode *>(baseNode)->GetNoSSANode());
  else
    return dynamic_cast<AddrofNode *>(baseNode);
}

uint32 MEDoNeighborOpt::getOffset(BaseNode *baseNode) {
  uint32 offset = 0;
  switch (baseNode->GetOpCode()) {
    case OP_iassign: {
      IassignNode *iassign = dynamic_cast<IassignNode *>(baseNode);
      MIRPtrType *ptrType =
          dynamic_cast<MIRPtrType *>(GlobalTables::GetTypeTable().GetTypeFromTyIdx(iassign->GetTyIdx()));
      MIRType *pointedType = ptrType->GetPointedType();
      if (pointedType->IsMIRStructType()) {
        offset = getStructOffset(iassign, ptrType, iassign->GetFieldID());
      } else {
        offset = getArrayOffset(iassign, ptrType, iassign->GetFieldID());
      }

      break;
    }
    case OP_iread: {
      IreadNode *iread = dynamic_cast<IreadNode *>(baseNode);
      CHECK_FATAL((iread != nullptr), "Invalid IR node with OpCode OP_iread.");
      MIRPtrType *ptrType =
          dynamic_cast<MIRPtrType *>(GlobalTables::GetTypeTable().GetTypeFromTyIdx(iread->GetTyIdx()));
      MIRType *pointedType = ptrType->GetPointedType();
      if (pointedType->IsMIRStructType()) {
        offset = getStructOffset(iread, ptrType, iread->GetFieldID());
      } else {
        offset = getArrayOffset(iread, ptrType, iread->GetFieldID());
      }
      break;
    }
    default: {
    }
  }

  return offset;
}

uint32 MEDoNeighborOpt::getStructOffset(BaseNode *node, MIRPtrType *ptrType, FieldID fieldID) {
  MIRType* pointedType = ptrType->GetPointedType();
  if (fieldID == 0) {
    uint32 size = pointedType->GetSize();
    this->node2size[node] = size;
    this->node2offset[node] = 0;
    return 0;
  } else {
    MIRStructType* tmp = dynamic_cast<MIRStructType*>(pointedType);
    CHECK_FATAL((tmp != nullptr), "Unknown type with fieldID access.");
    uint32 offset = tmp->GetFieldOffsetFromBaseAddr(fieldID).byteOffset;
    this->node2offset[node] = offset;
    MIRType *mirtype = GlobalTables::GetTypeTable().GetTypeFromTyIdx(tmp->GetFieldTyIdx(fieldID));
    this->node2size[node] = mirtype->GetSize();
    return offset;
  }
}

uint32 MEDoNeighborOpt::getArrayOffset(BaseNode *node, MIRPtrType *ptrType, FieldID fieldID) {
  BaseNode *opnd = node->Opnd(0);
  // opnd->Dump();
  if (opnd->GetOpCode() == OP_add) {
    // pattern: base + offset
    BaseNode *addrExpr = opnd->Opnd(0);
    BaseNode *offsetExpr = opnd->Opnd(1);
    if (addrExpr->GetOpCode() == OP_dread && offsetExpr->GetOpCode() == OP_constval) {
      ConstvalNode *offsetNode = dynamic_cast<ConstvalNode *>(offsetExpr);
      if (offsetNode == nullptr) return -1;
      MIRIntConst *intConst = dynamic_cast<MIRIntConst *>(offsetNode->GetConstVal());
      if (intConst == nullptr) return -1;
      uint32 offset = intConst->GetExtValue();

      MIRType *pointedType = ptrType->GetPointedType();
      size_t size = pointedType->GetSize();
      this->node2size[node] = size;
      this->node2offset[node] = offset;
      // LogInfo::MapleLogger() << "getArrayOffset\n";
      // node->Dump();
      // LogInfo::MapleLogger() << "size: " << size << " offset: " << offset << "\n";
      return offset;
    }
  }
  return 0xffffffff;
}

void MEDoNeighborOpt::DumpIntermediateResult() {
  std::map<const StmtNode *, std::vector<BaseNode *>>::iterator itr;
  for (itr = this->def2nodes.begin(); itr != this->def2nodes.end(); ++itr) {
    const StmtNode *defStmt = itr->first;
    std::vector<BaseNode *> targetNodes = itr->second;
    LogInfo::MapleLogger() << "defStmt:\n";
    defStmt->Dump();
    LogInfo::MapleLogger() << "use Nodes:\n";
    for (auto *useNode : targetNodes) {
      useNode->Dump();
      LogInfo::MapleLogger() << "size: " << this->node2size[useNode] << "\n";
      LogInfo::MapleLogger() << "offset: " << this->node2offset[useNode] << "\n";
    }
  }
}

void MEDoNeighborOpt::DumpResult() {
  std::vector<AsanNodePos>::iterator itr;
  for (itr = this->results.begin(); itr != this->results.end(); ++itr) {
    SrcPosition stmtPos = itr->stmtPos;
    std::vector<int> directions = itr->directions;
    Opcode op = itr->op;
    int64_t newSize = itr->newSize;

    LogInfo::MapleLogger() << "stmtPos:\n";
    stmtPos.Dump();
    LogInfo::MapleLogger() << "directions:\n";
    for (auto dir : directions) {
      LogInfo::MapleLogger() << dir << " ";
    }
    LogInfo::MapleLogger() << "\n";
    LogInfo::MapleLogger() << "op: " << int(op) << "\n";
    LogInfo::MapleLogger() << "newSize: " << newSize << "\n";
  }
}

std::vector<std::pair<BaseNode *, int>> MEDoNeighborOpt::sortOffset() {
  std::vector<std::pair<BaseNode *, int>> sortVec;
  for (auto &it : this->node2offset) {
    sortVec.push_back(it);
  }
  std::sort(sortVec.begin(), sortVec.end(), this->cmp);
  return sortVec;
}

void MEDoNeighborOpt::NeighborOptEntry(BB *bb) {
  this->node2offset.clear();
  this->node2size.clear();
  this->node2def.clear();
  this->def2nodes.clear();

  for (StmtNode &stmt : bb->GetStmtNodes()) {
    /*  cannot mark on current function directly,
        because [irmapbuild + premeemi] phases will rebuild (optimize) the function.
        have no choice but a StmtNode level granularity optimization
    */

    // visit all node, to find def stmt node
    std::vector<int> directions;
    TraverseNodes(&stmt, &stmt, directions);
  }

  // filter and update result
  filterRemovableMemAccesses();
}

void MEDoNeighborOpt::TraverseNodes(StmtNode *stmt, BaseNode *baseNode, std::vector<int> curDirections) {
  // LogInfo::MapleLogger() << "MEDoNeighborOpt::TraverseNodes \n";
  // stmt->Dump();
  // baseNode->Dump(2);
  // LogInfo::MapleLogger() << "\n";
  // SrcPosition srcPos = stmt->GetSrcPos();
  // srcPos.Dump();
  if (baseNode != nullptr) {
    BaseNode *mapKey = baseNode;
    if (baseNode->IsSSANode()) {
      mapKey = dynamic_cast<SSANode *>(baseNode)->GetNoSSANode();
    }
    this->node2stmt[mapKey] = stmt;
    this->node2directions[mapKey] = curDirections;

    switch (baseNode->GetOpCode()) {
      case OP_iassign: {
        IassignNode *iassign = dynamic_cast<IassignNode *>(baseNode);
        CHECK_FATAL((iassign != nullptr), "Invalid IR node with OpCode OP_iassign");
        BaseNode *addrExpr = iassign->Opnd(0);

        // struct pattern: we want to see where the addr-expr comes from
        auto *addrofSSANode = dynamic_cast<const AddrofSSANode *>(addrExpr);

        // array pattern: base ptr + offset
        if (addrExpr->GetOpCode() == OP_add) {
          addrExpr = addrExpr->Opnd(0);
          // we want to see where the addr-expr comes from
          addrofSSANode = dynamic_cast<const AddrofSSANode *>(addrExpr);
        }

        if (addrofSSANode == nullptr) break;  // without def-use info, or not an addrof node

        // Collect information
        CHECK_FATAL((addrofSSANode != nullptr), "not a valid SSA Node");
        const VersionSt *verSt = addrofSSANode->GetSSAVar();
        CHECK_FATAL((verSt != nullptr), "not a valid SSA Var");
        findDefStmt(iassign, *verSt);
        getOffset(iassign);

        break;
      }
      case OP_iassignoff: {
        IassignoffNode *iassignoff = dynamic_cast<IassignoffNode *>(baseNode);
        CHECK_FATAL((iassignoff != nullptr), "Invalid IR node with OpCode OP_iassignoff");
        // TODO
        // struct MemoryAccess memoryAccess = getIassignoffMemoryAccess(*iassignoff);

        // // the rhs-expr can still read from somewhere, push it to stack
        // baseNodeStack.push(iassignoff->GetBOpnd(1));
        break;
      }
      case OP_iassignfpoff:
      case OP_iassignpcoff:
        break;
      case OP_iread: {
        IreadNode *iread = nullptr;
        if (baseNode->IsSSANode()) {
          iread = dynamic_cast<IreadNode *>(dynamic_cast<IreadSSANode *>(baseNode)->GetNoSSANode());
        } else {
          iread = dynamic_cast<IreadNode *>(baseNode);
        }
        CHECK_FATAL((iread != nullptr), "Invalid IR node with OpCode OP_iread.");

        BaseNode *addrExpr = iread->Opnd(0);

        // struct pattern: we want to see where the addr-expr comes from
        auto *addrofSSANode = dynamic_cast<const AddrofSSANode *>(addrExpr);

        // pattern: base ptr + offset
        if (addrExpr->GetOpCode() == OP_add) {
          addrExpr = addrExpr->Opnd(0);
          // we want to see where the addr-expr comes from
          addrofSSANode = dynamic_cast<AddrofSSANode *>(addrExpr);
        }

        if (addrofSSANode == nullptr) break;  // without def-use info

        // Collect information
        const VersionSt *verSt = addrofSSANode->GetSSAVar();
        CHECK_FATAL((verSt != nullptr), "not a valid SSA Var");
        findDefStmt(iread, *verSt);
        getOffset(iread);

        break;
      }
      case OP_ireadoff:
      case OP_ireadfpoff:
      case OP_ireadpcoff:
        break;
      case OP_block: {
        CHECK_FATAL((false), "OP_block should not be visited in MEDoNeighborOpt");
        break;
      }
      default: {
      }
    }

    for (size_t j = 0; j < baseNode->NumOpnds(); ++j) {
      if (baseNode->GetOpCode() == OP_return) continue;
      std::vector<int> newDirections;
      std::copy(curDirections.begin(), curDirections.end(), back_inserter(newDirections));
      newDirections.push_back(j);

      TraverseNodes(stmt, baseNode->Opnd(j), newDirections);
    }
  }
}

const StmtNode *MEDoNeighborOpt::findDefStmt(BaseNode *baseNode, const VersionSt &vst) {
  // LogInfo::MapleLogger() << "Neighbor Opt: findDefStmt: ";
  // vst.Dump();
  // LogInfo::MapleLogger() << "\nNeighbor Opt: dump def stmt, type: ";

  const StmtNode *defStmt;
  if (vst.GetDefType() == VersionSt::kAssign) {
    // LogInfo::MapleLogger() << "kAssign\n";
    // vst.DumpDefStmt();
    const StmtNode *assign = vst.GetAssignNode();
    defStmt = assign;
  } else if (vst.GetDefType() == VersionSt::kPhi) {
    // LogInfo::MapleLogger() << "kPhi\n";
    // vst.DumpDefStmt();
    defStmt = nullptr;
    // PhiNode *phi = vst.GetPhi();
    // ASSERT(phi->GetResult() == &vst, "MarkVst: wrong corresponding version st in phi");
    // MarkControlDependenceLive(ToRef(dfBB));
    // for (size_t i = 0; i < phi->GetPhiOpnds().size(); ++i) {
    //   const VersionSt *verSt = phi->GetPhiOpnds()[i];
    //   AddToWorkList(verSt);
    // }
    // phi->SetIsLive(true);
  } else if (vst.GetDefType() == VersionSt::kMayDef) {
    // LogInfo::MapleLogger() << "kMayDef\n";
    // vst.DumpDefStmt();
    defStmt = nullptr;
    // const MayDefNode *mayDef = vst.GetMayDef();
    // ASSERT(mayDef->GetResult() == &vst, "MarkVst: wrong corresponding version st in maydef");
    // const VersionSt *verSt = mayDef->GetOpnd();

    // auto defStmt = mayDef->GetStmt();
    // if (kOpcodeInfo.IsCallAssigned(defStmt->GetOpCode())) {
    //   MapleVector<MustDefNode> &mustDefs = ssaTab.GetStmtMustDefNodes(*defStmt);
    //   for (auto &node : mustDefs) {
    //     if (aliasInfo->MayAlias(node.GetResult()->GetOst(), vst.GetOst())) {
    //       AddToWorkList(node.GetResult());
    //     }
    //   }
    // }

    // MarkStmtRequired(ToRef(defStmt), ToRef(dfBB));
    // AddToWorkList(verSt);
  } else if (vst.GetDefType() == VersionSt::kMustDef) {
    // LogInfo::MapleLogger() << "kMustDef\n";
    // vst.DumpDefStmt();
    const MustDefNode *mustDef = vst.GetMustDef();
    ASSERT(mustDef->GetResult() == &vst, "MarkVst: wrong corresponding version st in mustDef");
    const StmtNode *stmtNode = mustDef->GetStmt();
    defStmt = stmtNode;
    // MarkStmtRequired(ToRef(mustDef->GetStmt()), ToRef(dfBB));
    defStmt->Dump();
  } else {
    defStmt = nullptr;
  }
  this->node2def[baseNode] = defStmt;
  if (this->def2nodes.find(defStmt) == this->def2nodes.end()) {
    std::vector<BaseNode *> tmp;
    tmp.push_back(baseNode);
    this->def2nodes[defStmt] = tmp;
  } else {
    std::vector<BaseNode *> &tmp = this->def2nodes[defStmt];
    tmp.push_back(baseNode);
  }
  // LogInfo::MapleLogger() << "\n";
  return this->node2def[baseNode];
}

void MEDoNeighborOpt::filterRemovableMemAccesses() {
  std::map<const StmtNode *, std::vector<BaseNode *>>::iterator itr;
  for (itr = this->def2nodes.begin(); itr != this->def2nodes.end(); ++itr) {
    const StmtNode *defStmt = itr->first;
    if (defStmt == nullptr) continue;
    std::vector<BaseNode *> targetNodes = itr->second;

    // find removable nodes in targetNodes
    std::vector<std::pair<BaseNode *, int>> sortVec;
    for (BaseNode *node : targetNodes) {
      CHECK_FATAL((!node->IsSSANode()), "Find SSANode in NeighborOpt results.");

      sortVec.push_back(std::make_pair(node, this->node2offset[node]));
    }
    sort(sortVec.begin(), sortVec.end(), this->cmp);
    std::pair<BaseNode *, int> prevTarget = sortVec[0];

    int64_t newSize = 0;
    for (int i = 1; i < sortVec.size(); i++) {
      auto it = sortVec[i];
      // LogInfo::MapleLogger() << "offset2: " << it.second << " offset1: " << prevTarget.second << " size2: " << this->node2size[it.first] << "\n";
      if (it.second - prevTarget.second + this->node2size[it.first] <= 8) {
        // removable
        if (this->node2stmt.count(it.first) == 0) {
          LogInfo::MapleLogger() << "not in this->node2stmt \n";
          it.first->Dump();
        }
        StmtNode *stmtNode = dynamic_cast<StmtNode *>(this->node2stmt[it.first]);
        CHECK_FATAL((stmtNode != nullptr), "Not a valid StmtNode (1).");

        this->result.push_back(std::make_pair(stmtNode->GetSrcPos(), -1));
        AsanNodePos nodePos = {stmtNode->GetSrcPos(), this->node2directions[it.first], it.first->GetOpCode(), -1LL};
        this->results.push_back(nodePos);

        LogInfo::MapleLogger() << "removable:\n";
        it.first->Dump();
        stmtNode->Dump();
        LogInfo::MapleLogger() << "offset: " << it.second << "\n";
        LogInfo::MapleLogger() << "offset2: " << it.second << " offset1: " << prevTarget.second
                               << " size2: " << this->node2size[it.first] << "\n";
        newSize = it.second - prevTarget.second + this->node2size[it.first];
      } else {
        LogInfo::MapleLogger() << "new check:\n";
        StmtNode *stmtNode = dynamic_cast<StmtNode *>(this->node2stmt[prevTarget.first]);
        CHECK_FATAL((stmtNode != nullptr), "Not a valid StmtNode (2).");
        stmtNode->Dump();
        LogInfo::MapleLogger() << "new size: " << newSize << "\n";

        this->result.push_back(std::make_pair(stmtNode->GetSrcPos(), newSize));
        AsanNodePos nodePos = {stmtNode->GetSrcPos(), this->node2directions[prevTarget.first],
                               prevTarget.first->GetOpCode(), newSize};
        this->results.push_back(nodePos);
        prevTarget = it;
        newSize = 0;
      }
    }
    LogInfo::MapleLogger() << "new check:\n";
    StmtNode *stmtNode = dynamic_cast<StmtNode *>(this->node2stmt[prevTarget.first]);
    CHECK_FATAL((stmtNode != nullptr), "Not a valid StmtNode (3).");
    stmtNode->Dump();
    LogInfo::MapleLogger() << "new size: " << newSize << "\n";

    this->result.push_back(std::make_pair(stmtNode->GetSrcPos(), newSize));
    AsanNodePos nodePos = {stmtNode->GetSrcPos(), this->node2directions[prevTarget.first],
                           prevTarget.first->GetOpCode(), newSize};
    this->results.push_back(nodePos);
  }
}

void MEDoNeighborOpt::GetAnalysisDependence(maple::AnalysisDep &aDep) const {
  aDep.SetPreservedAll();
}

std::vector<AsanNodePos> MEDoNeighborOpt::GetResult() {
  // DumpResult(); // for debug
  return this->results;
}

bool MEDoNeighborOpt::PreAnalysis(MeFunction &mefunc) {
  // MIRBuilder *builder = mefunc.GetMIRModule().GetMIRBuilder();
  // this->func = &mefunc;
  if (mefunc.GetMirFunc()->GetAttr(FUNCATTR_extern)) {
    return false;
  }
  if (mefunc.GetName().find("__asan_") == 0 || mefunc.GetName().find("__san_cov_") == 0) {
    return false;
  }

  std::vector<StmtNode *> toInstrument;
  std::vector<StmtNode *> noReturnCalls;

  for (BB *bb : mefunc.GetCfg()->GetAllBBs()) {
    // bb could be nullptr sometime?
    if (bb == nullptr) continue;
    for (StmtNode &stmt : bb->GetStmtNodes()) {
      toInstrument.push_back(&stmt);
      // result.push_back(stmt.GetSrcPos()); // TODO
      if (CallNode *callNode = dynamic_cast<CallNode *>(&stmt)) {
        MIRFunction *calleeFunc = GlobalTables::GetFunctionTable().GetFunctionFromPuidx(callNode->GetPUIdx());
        if (calleeFunc->NeverReturns() || calleeFunc->GetName() == "exit") {
          noReturnCalls.push_back(callNode);
        }
      }
    }
  }

  // toInstrument contains basically all statements in a function
  // we want to do Neighbor Checks optimization before further instrumentation
  /*/ TODO: Monkbai 6/27/2023
      1. Collect base addresses for all indirect memory access
      2. Collect offset for all indirect memory access
      3. Find removable checks and update toInstrument
  /*/
  MeCFG *cfg = mefunc.GetCfg();
  // MIRModule &mod = mefunc.GetMIRModule();
  for (BB *bb : cfg->GetAllBBs()) {
    /*
    LogInfo::MapleLogger() << "BB:\n"; 
    bb->Dump(&mod); 
    LogInfo::MapleLogger() << "\n";
    */

    // for each BB, we check all stmt inside
    if (bb == nullptr) continue;
    NeighborOptEntry(bb);
  }

  return true;
}

bool MEDoNeighborOpt::PhaseRun(MeFunction &f) {
  LogInfo::MapleLogger() << "The MEDoNeighborOpt::PhaseRun is running " << f.GetName() << "\n";
  PreAnalysis(f);
  // for (auto pos : this->result){
  //   LogInfo::MapleLogger() << "filenum " << pos.FileNum() << " linenum " << pos.LineNum() << " colnum " << pos.Column() << "\n";
  // }
  LogInfo::MapleLogger() << "The MEDoNeighborOpt::PhaseRun ends " << f.GetName() << "\n";
  return true;
}

}  // namespace maple

#endif